import torch.nn as nn

"""
前馈网络层
"""


class FeedForwardNet(nn.Module):
    def __init__(self, d_model=512, d_ff=2048):
        super(FeedForwardNet, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.activation = nn.ReLU()
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.linear2(self.activation(self.linear1(x)))


if __name__ == '__main__':
    import torch
    from embedding_layer import EmbeddingLayer

    test_input = torch.randint(0, 10, (4, 8))
    embedding = EmbeddingLayer(10, 512)
    test_input = embedding(test_input)
    feed_forward = FeedForwardNet()
    output = feed_forward(test_input)
    print(output.shape)
    print(output)

    from add_and_norm_layer import AddAndNorm

    add_and_norm = AddAndNorm()
    output = add_and_norm(test_input, feed_forward)
    print(output.shape)
    print(output)
